#undef NDEBUG
#include <lcthw/hashmap.h>
#include <lcthw/dbg.h>
#include <lcthw/bstrlib.h>

static int default_compare(void *a, void *b)
{
    return bstrcmp((bstring)a, (bstring)b);
}

/**
 * Simple Bob Jenkins's hash algorithm taken from the wikipedia description.
 */
static uint32_t default_hash(void *a)
{
    size_t len = blength((bstring)a);
    char *key = bdata((bstring)a); // get the data from the bstring, cast it to char pointer
    uint32_t hash = 0;
    uint32_t i = 0;

    for (hash = i = 0; i < len; ++i)
    {
        hash += key[i];
        hash += (hash << 10);
        hash ^= (hash >> 6);
    }

    hash += (hash << 3);
    hash ^= (hash >> 11);
    hash += (hash << 15);

    return hash;
}

// Hashmap_create: Creates a new hashmap, returns a pointer to the hashmap.
Hashmap *Hashmap_create(Hashmap_compare compare, Hashmap_hash hash)
{
    Hashmap *map = calloc(1, sizeof(Hashmap)); // allocate memory for the hashmap
    check_mem(map);

    map->compare = compare == NULL ? default_compare : compare;                // set the compare function pointer
    map->hash = hash == NULL ? default_hash : hash;                            // set the hash function pointer
    map->buckets = DArray_create(sizeof(DArray *), DEFAULT_NUMBER_OF_BUCKETS); // create a DArray of buckets
    map->buckets->end = map->buckets->max;                                     // set the end of the DArray to the max

    return map;

error:
    if (map)
    {
        Hashmap_destroy(map);
    }

    return NULL;
}

// Hashmap_destroy: Destroys the hashmap
void Hashmap_destroy(Hashmap *map)
{
    int i = 0;
    int j = 0;

    if (map)
    {
        if (map->buckets)
        {
            for (i = 0; i < DArray_count(map->buckets); i++)
            {
                DArray *bucket = DArray_get(map->buckets, i);
                if (bucket)
                {
                    for (j = 0; j < DArray_count(bucket); j++)
                    {
                        free(DArray_get(bucket, j));
                    }
                    DArray_destroy(bucket);
                }
            }
            DArray_destroy(map->buckets);
        }

        free(map);
    }
}

// Hashmap_node_create: Creates a new hashmap node
// Inputs: hash - hash value of the key; key - pointer to the key; data - pointer to the data
// Returns: pointer to the hashmap node
static inline HashmapNode *Hashmap_node_create(int hash, void *key, void *data)     // static inline functions are used to define functions that are only used in one file
{
    HashmapNode *node = calloc(1, sizeof(HashmapNode));
    check_mem(node);

    node->key = key;
    node->data = data;
    node->hash = hash;

    return node;    

error:
    return NULL;
}

// hashmap_find_node: Finds the node in the hashmap
// Inputs: map - pointer to the hashmap; key - pointer to the key; hash_out - pointer to the hash value; create - flag to create a new node
// Returns: pointer to the hashmap node
static inline DArray *Hashmap_find_bucket(Hashmap *map, void *key, int create, uint32_t *hash_out)
{
    uint32_t hash = map->hash(key); // generate the hash value of the key
    int bucket_n = hash % DEFAULT_NUMBER_OF_BUCKETS; // get the bucket number
    check(bucket_n >= 0, "Invalid bucket found: %d", bucket_n); // check if the bucket number is valid

    // store the hash value in the hash_out pointer, so that it can be used outside the function
    *hash_out = hash; // set the hash_out to the hash value

    DArray *bucket = DArray_get(map->buckets, bucket_n); // get the bucket from the DArray

    if (!bucket && create) // if the bucket is not found and create is true
    {
        bucket = DArray_create(sizeof(void *), DEFAULT_NUMBER_OF_BUCKETS); // create a new bucket
        check_mem(bucket);
        DArray_set(map->buckets, bucket_n, bucket); // set the bucket in the DArray
    }

    return bucket;

error:
    return NULL;
}

// Hashmap_set: Adds a key-value pair to the hashmap
// Inputs: map - pointer to the hashmap; key - pointer to the key; data - pointer to the data
// Returns: 0 if successful, -1 if an error occurs
int Hashmap_set(Hashmap *map, void *key, void *data)
{
    uint32_t hash = 0;
    DArray *bucket = Hashmap_find_bucket(map, key, 1, &hash); // find the bucket in the hashmap
    check(bucket, "Error can't create bucket.");

    HashmapNode *node = Hashmap_node_create(hash, key, data); // create a new hashmap node
    check_mem(node);

    DArray_push(bucket, node); // push the node to the bucket

    return 0;

error:
    return -1;
}

// hashmap_get_node: Gets the node from the bucket of the hashmap
// Inputs: map - pointer to the hashmap; key - pointer to the key; hash - hash value of the key; bucket - pointer to the bucket
// Returns: index of the node in the bucket
static inline int Hashmap_get_node(Hashmap *map, void *key, uint32_t hash, DArray *bucket)
{
    int i = 0;

    for (i = 0; i < DArray_end(bucket); i++) // iterate through the bucket
    {
        debug("TRY: %d", i);
        HashmapNode *node = DArray_get(bucket, i); // get the node from the bucket
        if (node->hash == hash && map->compare(node->key, key) == 0) // if the hash value and key match
        {
            return i; // return the index
        }
    }

    return -1;
}

// Hashmap_get: Returns the data associated with the key
void *Hashmap_get(Hashmap *map, void *key)
{
    uint32_t hash = 0;
    DArray *bucket = Hashmap_find_bucket(map, key, 0, &hash); // find the bucket in the hashmap, do not create a new bucket
    if (!bucket) // if the bucket is not found
    {
        return NULL;
    }

    int i = Hashmap_get_node(map, key, hash, bucket); // get the index of the node in the bucket
    if (i == -1) // if the node is not found
    {
        return NULL;
    }

    HashmapNode *node = DArray_get(bucket, i); // get the node from the bucket
    check(node != NULL, "Failed to get node from bucket when it should exist.");

    return node->data; // return the data associated with the key

error: // error handling
    return NULL;
}

// Hashmap_traverse: Traverses the hashmap
int Hashmap_traverse(Hashmap *map, Hashmap_traverse_cb traverse_cb)
{
    int i = 0;
    int j = 0;
    int rc = 0;

    for (i = 0; i < DArray_count(map->buckets); i++) // iterate through the buckets
    {
        DArray *bucket = DArray_get(map->buckets, i); // get the bucket from the DArray
        if (bucket) // if the bucket is not NULL
        {
            for (j = 0; j < DArray_count(bucket); j++) // iterate through the bucket
            {
                HashmapNode *node = DArray_get(bucket, j); // get the node from the bucket
                rc = traverse_cb(node); // call the traverse callback function
                if (rc != 0) // if the callback function returns an error
                {
                    return rc;  // return the error code
                }
            }
        }
    }

    return 0;
}

// Hashmap_delete: Deletes the key-value pair from the hashmap
void *Hashmap_delete(Hashmap *map, void *key)
{
    uint32_t hash = 0;
    DArray *bucket = Hashmap_find_bucket(map, key, 0, &hash); // find the bucket in the hashmap, do not create a new bucket
    if (!bucket) // if the bucket is not found
    {
        return NULL;
    }

    int i = Hashmap_get_node(map, key, hash, bucket); // get the index of the node in the bucket
    if (i == -1) // if the node is not found
    {
        return NULL;
    }

    HashmapNode *node = DArray_get(bucket, i); // get the node from the bucket
    void *data = node->data; // get the data from the node
    free(node); // free the node
    HashmapNode *ending = DArray_pop(bucket); // pop the last node from the bucket

    if (ending != node) // if the ending node is not the same as the node
    {
        // alright looks like it's not the last one, swap it
        // use the original ending node to replace the deleted node
        DArray_set(bucket, i, ending); // set the ending node in the bucket
    }

    return data;
}